from collections import defaultdict
from struct import unpack

_PKT_HEADER = 0xBE11E2
_PKT_FOOTER = 0x62504944

MASKS = {'ASIC_ROW':   0x30000000,
         'ASIC_COL':   0x03000000,
         'PACKET_NO':  0x00FF0000,
         'ASIC_CHAN':  0x00007000,         
         'ORIGIN_WIN': 0x000001FF,
         }


def to_words(bytestr):
    if len(bytestr) % 4 != 0:
        msg = "Length of data (%d) is not a multiple of 4"
        raise ValueError(msg % len(bytestr))
    
    nwords = len(bytestr) / 4
    return unpack("%dI" % nwords, bytestr)


class PacketStream(object):
    packet_length = 140

    def __init__(self, fp):
        if type(fp) == type("hello"):
            self.fp = open(fp)
        else:
            self.fp = fp

        self._pktbuf = []
        self._bufsize_pkts = 100
        self._bufsize_words = self.packet_length * self._bufsize_pkts
        self._bufsize_bytes = self._bufsize_words * 4

    def __iter__(self):
        return self

    def next(self):
        if len(self._pktbuf) == 0:
            self._load_pktbuf()

        try:
            return self._pktbuf.pop(0) # Use fixed-size numpy buffer here
        except IndexError:
            raise StopIteration
                
    def reset(self):
        self.fp.seek(0)
        
    def _load_pktbuf(self):
        self._pktbuf = []

        bytestr = self.fp.read(self._bufsize_bytes)

        if len(bytestr) == 0:
            return 

        words = to_words(bytestr)

        if len(words) % 140 != 0:
            msg = "Read %d words, which is not a multiple of packet length (%d)"
            raise ValueError(msg % (len(words), self.packet_length))

        if len(words) != self._bufsize_words:
            npkts = len(words) / self.packet_length
        else:
            npkts = self._bufsize_pkts
        
        
        for i in xrange(npkts - 1):
            start = i * self.packet_length
            stop  = (i+1) * self.packet_length

            p = words[start:stop]
            
            if len(p) < self.packet_length:
                msg = "Read less than the expected packet length " + \
                    "(%d read, %d expected)"
                raise ValueError(msg % (len(p), self.packet_length))
                
            self._pktbuf.append(p)

        return


class Waveform(object):
    def __init__(self, row, col, chan, win, samps):
        self.row     = row
        self.col     = col
        self.chan    = chan
        self.win     = win
        self.samples = samps[:]

        self.iread    = 0 # Which window read is this, 0, 1, 2, or 3?
    
    def __iter__(self):
        return iter(self.samples)
    


def wfs_from_packet(pkt):
    wfs = []
    for i in xrange(4):
        base = 5 + 33*i
        origin = pkt[base]

        row   = (origin & MASKS['ASIC_ROW']) >> 28
        col   = (origin & MASKS['ASIC_COL']) >> 24
        chan  = (origin & MASKS['ASIC_CHAN']) >> 12
        win   = origin & MASKS['ORIGIN_WIN']

        samps = []
        for j in xrange(base + 1, base + 33):
            adc0 = (pkt[j] & 0xFFFF0000) >> 16
            adc1 = (pkt[j] & 0xFFFF) >> 16

            samps.append(adc0)
            samps.append(adc1)

        wfs.append(Waveform(row, col, chan, win, samps))

    return wfs


class Event(object):
    def __init__(self):
        self.evtno = 0
        self.wfs = []

        self.idmap = defaultdict(lambda: 0)

    def add_header(self, pkt):
        self.evtno = pkt[4]

    def add_wf(self, pkt):
        for wf in wfs_from_packet(pkt):
            wfid = "%d%d%d" % (wf.row, wf.col, wf.chan)

            wf.iread = self.idmap[wfid]
            self.idmap[wfid] += 1

            self.wfs.append(wf)

    def __iter__(self):
        return iter(self.wfs)


class EventStream(object):
    def __init__(self, fp):
        self.ps = PacketStream(fp)
    
    def __iter__(self):
        return self

    def next(self):
        evt = Event()
        wfps = []

        pkt = self.ps.next()
        while pkt[3] != 0xEADA:
            pkt = self.ps.next()

        evt.add_header(pkt)

        pkt = self.ps.next()
        while pkt[3] == 0xC0FFEE:
            evt.add_wf(pkt)
            pkt = self.ps.next()
                                
        return evt
        

        
